from flask import Blueprint, request, jsonify
from app.extensions import db
from app.models.budgets import Budget, CategoryBudget, BudgetType
from app.models.transactions import Transaction
from app.models.categories import Category
from app.routes.auth import token_required
from datetime import datetime

budget_bp = Blueprint('budget', __name__)

def calculate_spent_amount(user_id, year, month=None):
    """Calculate total spent amount for given period"""
    query = db.session.query(db.func.sum(Transaction.amount)).filter(
        Transaction.user_id == user_id,
        Transaction.type == 'expense',
        db.extract('year', Transaction.date) == year
    )
    
    if month:
        query = query.filter(db.extract('month', Transaction.date) == month)
    
    result = query.scalar()
    return float(result) if result else 0.0

def calculate_category_spent(category_id, user_id, year, month=None):
    """Calculate spent amount for specific category"""
    query = db.session.query(db.func.sum(Transaction.amount)).filter(
        Transaction.user_id == user_id,
        Transaction.category_id == category_id,
        Transaction.type == 'expense',
        db.extract('year', Transaction.date) == year
    )
    
    if month:
        query = query.filter(db.extract('month', Transaction.date) == month)
    
    result = query.scalar()
    return float(result) if result else 0.0

@budget_bp.route('/api/budgets/total', methods=['GET'])
@token_required
def get_total_budget(current_user):
    """获取总预算及支出接口"""
    try:
        type_str = request.args.get('type')
        year = int(request.args.get('year'))
        month = int(request.args.get('month')) if request.args.get('month') else None
        
        # 将字符串转换为枚举
        try:
            type_enum = BudgetType(type_str)
        except ValueError:
            return jsonify({'success': False, 'message': '预算类型必须是 month 或 year'}), 400
        
        if type_enum == BudgetType.MONTH and not month:
            return jsonify({'success': False, 'message': '月预算必须提供月份'}), 400
        
        # Get budget
        budget = Budget.query.filter_by(
            user_id=current_user.user_id,
            type=type_enum,
            year=year,
            month=month
        ).first()
        
        # Calculate spent amount
        total_spent = calculate_spent_amount(current_user.user_id, year, month)
        
        response_data = {
            'totalBudget': float(budget.total_budget) if budget else 0.0,
            'totalSpent': total_spent,
            'year': year,
            'type': type_str
        }
        
        if month:
            response_data['month'] = month
        
        return jsonify({
            'success': True,
            'data': response_data
        })
        
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/budgets/total', methods=['PUT'])
@token_required
def update_total_budget(current_user):
    """更新总预算接口"""
    try:
        data = request.get_json()
        type_str = data.get('type')
        year = int(data.get('year'))
        month = int(data.get('month')) if data.get('month') else None
        amount = float(data.get('amount'))
        
        try:
            type_enum = BudgetType(type_str)
        except ValueError:
            return jsonify({'success': False, 'message': '预算类型必须是 month 或 year'}), 400
        
        if type_enum == BudgetType.MONTH and not month:
            return jsonify({'success': False, 'message': '月预算必须提供月份'}), 400
        
        if amount <= 0:
            return jsonify({'success': False, 'message': '预算金额必须大于0'}), 400
        
        # Find existing budget or create new
        budget = Budget.query.filter_by(
            user_id=current_user.user_id,
            type=type_enum,
            year=year,
            month=month
        ).first()
        
        if budget:
            budget.total_budget = amount
            budget.updated_at = datetime.utcnow()
        else:
            budget = Budget(
                user_id=current_user.user_id,
                type=type_enum,
                year=year,
                month=month,
                total_budget=amount,
                created_at=datetime.utcnow(),
                updated_at=datetime.utcnow()
            )
            db.session.add(budget)
        
        db.session.commit()
        
        response_data = {
            'totalBudget': float(budget.total_budget),
            'year': year,
            'type': type_str
        }
        
        if month:
            response_data['month'] = month
        
        return jsonify({
            'success': True,
            'data': response_data
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/budgets/categories', methods=['GET'])
@token_required
def get_category_budgets(current_user):
    """获取分类预算列表接口"""
    try:
        type_str = request.args.get('type')
        year = int(request.args.get('year'))
        month = int(request.args.get('month')) if request.args.get('month') else None
        
        try:
            type_enum = BudgetType(type_str)
        except ValueError:
            return jsonify({'success': False, 'message': '预算类型必须是 month 或 year'}), 400
        
        if type_enum == BudgetType.MONTH and not month:
            return jsonify({'success': False, 'message': '月预算必须提供月份'}), 400
        
        # Get category budgets
        category_budgets = CategoryBudget.query.filter_by(
            user_id=current_user.user_id,
            type=type_enum,
            year=year,
            month=month
        ).all()
        
        result = []
        for cb in category_budgets:
            spent = calculate_category_spent(cb.category_id, current_user.user_id, year, month)
            
            result.append({
                'id': cb.id,
                'categoryId': cb.category_id,
                'categoryName': cb.category.name if cb.category else '未知分类',
                'budget': float(cb.amount),
                'spent': spent,
                'remaining': float(cb.amount) - spent,
                'createdAt': cb.created_at.isoformat() + 'Z'
            })
        
        response_data = {
            'categoryBudgets': result,
            'year': year,
            'type': type_str
        }
        
        if month:
            response_data['month'] = month
        
        return jsonify({
            'success': True,
            'data': response_data
        })
        
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/budgets/categories', methods=['POST'])
@token_required
def add_category_budget(current_user):
    """添加分类预算接口"""
    try:
        data = request.get_json()
        type_str = data.get('type')
        year = int(data.get('year'))
        month = int(data.get('month')) if data.get('month') else None
        category_id = int(data.get('categoryId'))
        budget_amount = float(data.get('budget'))
        
        try:
            type_enum = BudgetType(type_str)
        except ValueError:
            return jsonify({'success': False, 'message': '预算类型必须是 month 或 year'}), 400
        
        if type_enum == BudgetType.MONTH and not month:
            return jsonify({'success': False, 'message': '月预算必须提供月份'}), 400
        
        if budget_amount <= 0:
            return jsonify({'success': False, 'message': '预算金额必须大于0'}), 400
        
        # Check if category exists
        category = Category.query.get(category_id)
        if not category:
            return jsonify({'success': False, 'message': '分类不存在'}), 404
        
        # Check if budget already exists
        existing_budget = CategoryBudget.query.filter_by(
            user_id=current_user.user_id,
            category_id=category_id,
            type=type_enum,
            year=year,
            month=month
        ).first()
        
        if existing_budget:
            return jsonify({'success': False, 'message': '该分类在此时间段内已有预算'}), 400
        
        # Create new category budget
        category_budget = CategoryBudget(
            user_id=current_user.user_id,
            category_id=category_id,
            type=type_enum,
            year=year,
            month=month,
            amount=budget_amount
        )
        
        db.session.add(category_budget)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': category_budget.id,
                'categoryId': category_budget.category_id,
                'budget': float(category_budget.amount)
            }
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/budgets/categories/<int:id>', methods=['PUT'])
@token_required
def update_category_budget(current_user, id):
    """更新分类预算接口"""
    try:
        data = request.get_json()
        new_budget = float(data.get('budget'))
        
        if new_budget <= 0:
            return jsonify({'success': False, 'message': '预算金额必须大于0'}), 400
        
        category_budget = CategoryBudget.query.filter_by(
            id=id,
            user_id=current_user.user_id
        ).first()
        
        if not category_budget:
            return jsonify({'success': False, 'message': '分类预算不存在'}), 404
        
        category_budget.amount = new_budget
        category_budget.updated_at = datetime.utcnow()

        new_category_id = data.get('categoryId')
        if new_category_id is not None:
            # 假设 category_id 在数据库中是整数类型
            category_budget.category_id = int(new_category_id)  # 更新分类ID字段

        category_budget.updated_at = datetime.utcnow()
        
        db.session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': category_budget.id,
                'budget': float(category_budget.amount)
            }
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/budgets/categories/<int:id>', methods=['DELETE'])
@token_required
def delete_category_budget(current_user, id):
    """删除分类预算接口"""
    try:
        category_budget = CategoryBudget.query.filter_by(
            id=id,
            user_id=current_user.user_id
        ).first()
        
        if not category_budget:
            return jsonify({'success': False, 'message': '分类预算不存在'}), 404
        
        db.session.delete(category_budget)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': id
            }
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500

@budget_bp.route('/api/expenses/categories', methods=['GET'])
@token_required
def get_category_expenses(current_user):
    """获取分类支出数据接口"""
    try:
        type_str = request.args.get('type')
        year = int(request.args.get('year'))
        month = int(request.args.get('month')) if request.args.get('month') else None
        
        try:
            type_enum = BudgetType(type_str)
        except ValueError:
            return jsonify({'success': False, 'message': '类型必须是 month 或 year'}), 400
        
        if type_enum == BudgetType.MONTH and not month:
            return jsonify({'success': False, 'message': '月类型必须提供月份'}), 400
        
        # Get all categories
        categories = Category.query.all()
        
        result = []
        for category in categories:
            spent = calculate_category_spent(category.id, current_user.user_id, year, month)

            if spent > 0:  # Only include categories with expenses
                # Count transactions
                query = db.session.query(Transaction).filter(
                    Transaction.user_id == current_user.user_id,
                    Transaction.category_id == category.id,
                    Transaction.type == 'expense',
                    db.extract('year', Transaction.date) == year
                )

                if month:
                    query = query.filter(db.extract('month', Transaction.date) == month)

                expense_count = query.count()

                result.append({
                    'categoryId': category.id,
                    'categoryName': category.name,
                    'spent': spent,
                    'expenseCount': expense_count
                })
        
        response_data = {
            'categoryExpenses': result,
            'year': year,
            'type': type_str
        }
        
        if month:
            response_data['month'] = month
        
        return jsonify({
            'success': True,
            'data': response_data
        })
        
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500